{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Burr RAG with LanceDB and dlt document ingestion\n",
    "\n",
    "This example shows how to build a chatbot with RAG over Substack blogs (or any RSS feed) stored into LanceDB. \n",
    "\n",
    "The stack includes:\n",
    "\n",
    "- Burr: define your RAG logic \n",
    "- LanceDB: store and retrieve documents using vector search\n",
    "- dlt: ingest unstructured web pages and store them as structured documents\n",
    "- OpenAI: embed documents for vector search and generate answers using LLMs\n",
    "- OpenTelemetry: automatically track telemetry from Burr, LanceDB, and OpenAI in a unified way"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. `Substack -> LanceDB` ingestion with `dlt`\n",
    "\n",
    "To ingest data, we use [dlt and its LanceDB integration](https://dlthub.com/devel/dlt-ecosystem/destinations/lancedb), which makes it very simple to query, embed, and store blogs from the web into LanceDB tables. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Text processing\n",
    "\n",
    "First, we define simple functions to split long text strings into sentences and a way to assemble sentences into larger context windows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "def split_text(text):\n",
    "    \"\"\"Split text on punction (., !, ?).\"\"\"\n",
    "    sentence_endings = r'[.!?]+'\n",
    "    for sentence in re.split(sentence_endings, text):\n",
    "        sentence = sentence.strip()\n",
    "        if sentence:\n",
    "            yield sentence\n",
    "\n",
    "\n",
    "def contextualize(chunks: list[str], window=5, stride=3, min_window_size=2):\n",
    "    \"\"\"Rolling window operation to join consecutive sentences into larger chunks.\"\"\"\n",
    "    n_chunks = len(chunks)\n",
    "    for start_i in range(0, n_chunks, stride):\n",
    "        if (start_i + window <= n_chunks) or (n_chunks - start_i >= min_window_size):\n",
    "            yield \" \".join(chunks[start_i : min(start_i + window, n_chunks)])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Define `dlt` resources\n",
    "\n",
    "To use `dlt`, you author `Resource` objects that generate data using the `@dlt.resource` decorator. In this case, we create a resource that pulls an RSS feed from a Substack blog URL using the `requests` and `feedparser` libraries. Then, we iterate over RSS entries and yield them as dictionaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Generator\n",
    "\n",
    "import requests\n",
    "import feedparser\n",
    "import dlt\n",
    "\n",
    "@dlt.resource(name=\"substack\", write_disposition=\"merge\", primary_key=\"id\")\n",
    "def rss_entries(substack_url: str) -> Generator:\n",
    "    \"\"\"Substack blog entries retrieved from a RSS feed\"\"\"\n",
    "    FIELDS_TO_EXCLUDE = [\n",
    "        \"published_parsed\",\n",
    "        \"title_detail\",\n",
    "        \"summary_detail\",\n",
    "        \"author_detail\",\n",
    "        \"guidislink\",\n",
    "        \"authors\",\n",
    "        \"links\"\n",
    "    ]\n",
    "\n",
    "    r = requests.get(f\"{substack_url}/feed\")\n",
    "    rss_feed = feedparser.parse(r.content)\n",
    "    for entry in rss_feed[\"entries\"]:\n",
    "        for field in FIELDS_TO_EXCLUDE:\n",
    "            entry.pop(field)\n",
    "\n",
    "        yield entry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can then use `@dlt.transformer` to define operations on the values returned by `Resource` objects. In this case, we define three transformations that we'll chain:\n",
    "\n",
    "1. Parse HTML into a string stripped of tags\n",
    "2. Chunk the text string by splitting it into sentences\n",
    "3. Join sentence chunks into larger \"context windows\" via a rolling operation.\n",
    "\n",
    "We use a custom trick to map and store the relationship between HTML pages, sentence chunks, and context windows ([learn more](https://github.com/dlt-hub/dlt/issues/1699))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bs4 import BeautifulSoup\n",
    "import utils\n",
    "\n",
    "\n",
    "@dlt.transformer(primary_key=\"id\")\n",
    "def parsed_html(rss_entry: dict):\n",
    "    \"\"\"Parse the HTML from the RSS entry\"\"\"\n",
    "    soup = BeautifulSoup(rss_entry[\"content\"][0][\"value\"], \"html.parser\")\n",
    "    parsed_text = soup.get_text(separator=\" \", strip=True)\n",
    "    yield {\"id\": rss_entry[\"id\"], \"text\": parsed_text}\n",
    "\n",
    "\n",
    "@dlt.transformer(primary_key=\"chunk_id\")\n",
    "def chunks(parsed_html: dict) -> list[dict]:\n",
    "    \"\"\"Chunk text\"\"\"\n",
    "    return [\n",
    "        dict(\n",
    "            document_id=parsed_html[\"id\"],\n",
    "            chunk_id=idx,\n",
    "            text=text,\n",
    "        )\n",
    "        for idx, text in enumerate(split_text(parsed_html[\"text\"]))\n",
    "    ]\n",
    "\n",
    "# order is important for reduce / rolling step\n",
    "# default to order of the batch or specifying sorting key\n",
    "@dlt.transformer(primary_key=\"context_id\")\n",
    "def contexts(chunks: list[dict]) -> Generator:\n",
    "    \"\"\"Assemble consecutive chunks into larger context windows\"\"\"\n",
    "    # first handle the m-to-n relationship\n",
    "    # set of foreign keys (i.e., \"chunk_id\")\n",
    "    chunk_id_set = set(chunk[\"chunk_id\"] for chunk in chunks)\n",
    "    context_id = utils.hash_set(chunk_id_set)\n",
    "    \n",
    "    # create a table only containing the keys\n",
    "    for chunk_id in chunk_id_set :\n",
    "        yield dlt.mark.with_table_name(\n",
    "            {\"chunk_id\": chunk_id, \"context_id\": context_id},\n",
    "            \"chunks_to_contexts_keys\",\n",
    "        ) \n",
    "    \n",
    "    # main transformation logic\n",
    "    for contextualized in contextualize([chunk[\"text\"] for chunk in chunks]):\n",
    "        yield dlt.mark.with_table_name(\n",
    "            {\"context_id\": context_id, \"text\": contextualized},\n",
    "            \"contexts\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Execute the pipeline\n",
    "\n",
    "Before ingesting data, we need to the configuration for the `dlt` destination (LanceDB in our case). We specify which OpenAI model we want to use for text embedding and store our API key.\n",
    "\n",
    "`dlt` provides [multiple ways to do so](https://dlthub.com/devel/general-usage/credentials), but using the `os` module is simply the most convenient for this tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# set your OpenAI API key\n",
    "openai_api_key = ...\n",
    "\n",
    "# this environment variable isn't needed by dlt, but we'll use it later\n",
    "os.environ[\"OPENAI_API_KEY\"] = openai_api_key\n",
    "\n",
    "os.environ[\"DESTINATION__LANCEDB__EMBEDDING_MODEL_PROVIDER\"] = \"openai\"\n",
    "os.environ[\"DESTINATION__LANCEDB__EMBEDDING_MODEL\"] = \"text-embedding-3-small\"\n",
    "\n",
    "os.environ[\"DESTINATION__LANCEDB__CREDENTIALS__URI\"] = \".lancedb\"\n",
    "os.environ[\"DESTINATION__LANCEDB__CREDENTIALS__EMBEDDING_MODEL_PROVIDER_API_KEY\"] = openai_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we combine our Substack blog `Resource` with the different text processing `Transformer` using the pipe operator `|`. We also use the `lancedb_adapter` by `dlt` which allows to specify which field will be embed with the OpenAI embedding service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dlt.destinations.adapters import lancedb_adapter\n",
    "import dlt.destinations.impl.lancedb.models\n",
    "\n",
    "blog_url = \"https://blog.dagworks.io/\"\n",
    "\n",
    "full_entries = lancedb_adapter(rss_entries(blog_url), embed=\"summary\")\n",
    "chunked_entries = rss_entries(blog_url) | parsed_html | chunks\n",
    "contextualized_chunks = lancedb_adapter(chunked_entries | contexts, embed=\"text\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we create the `Pipeline` object and minimally specify a `pipeline_name` and `destination`. This won't exectue any code or ingest any data.\n",
    "\n",
    "Then, calling `pipeline.run()` with the `Resource` and `Transformer` objects will launch the ingestion job and return a `LoadInfo` object detailing the results of the job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = dlt.pipeline(\n",
    "    pipeline_name=\"substack-blog\",\n",
    "    destination=\"lancedb\",\n",
    "    dataset_name=\"dagworks\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pipeline substack-blog load step completed in 3.02 seconds\n",
      "1 load package(s) were loaded to destination LanceDB and into dataset dagworks\n",
      "The LanceDB destination used <dlt.destinations.impl.lancedb.configuration.LanceDBCredentials object at 0x7fbe84da5810> location to store data\n",
      "Load package 1724941288.3716326 is LOADED and contains no failed jobs\n"
     ]
    }
   ],
   "source": [
    "load_info = pipeline.run([full_entries, chunked_entries, contextualized_chunks])\n",
    "print(load_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Burr RAG with LanceDB memory\n",
    "Burr allows you to define an `Application` by defining a set of actions and valid transitions between them. This approach allows to define complex agents in an easy-to-understand and debug manner. \n",
    "\n",
    "Burr solves many challenges to productionize agents including monitoring, storing interactions, streaming, and more, and comes with a rich open-source UI for observability."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Define `@action`\n",
    "\n",
    "First, we define actions the agent can take with the `@action` decorator. The function must take a `State` object as first argument and return a `State` object. The decorator specifies which `State` fields can be read from and written to.\n",
    "\n",
    "The next cell contains two actions:\n",
    "- `relevant_chunk_retrieval()` reads from the LanceDB table `dagworks___contexts` that was generated by `dlt` and retrieves the top 4 most similar rows to the `user_query` string. It writes the search results to the `relevant_chunks` state field and appends the user input to the `chat_history` state field.\n",
    "- `bot_turn()` reads the `chat_history` and the `relevant_chunks` from state, combine the text into a prompt and send a request to OpenAI. The LLM's response is appended to the `chat_history` state field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import textwrap\n",
    "\n",
    "import openai\n",
    "import lancedb\n",
    "\n",
    "from burr.core import State, action\n",
    "\n",
    "\n",
    "@action(reads=[], writes=[\"relevant_chunks\", \"chat_history\"])\n",
    "def relevant_chunk_retrieval(\n",
    "    state: State,\n",
    "    user_query: str,\n",
    "    lancedb_con: lancedb.DBConnection,\n",
    ") -> State:\n",
    "    \"\"\"Search LanceDB with the user query and return the top 4 results\"\"\"\n",
    "    # this is a table generated by `dlt`\n",
    "    text_chunks_table = lancedb_con.open_table(\"dagworks___contexts\")\n",
    "\n",
    "    search_results = (\n",
    "        text_chunks_table\n",
    "        .search(user_query)  # this automatically embed the query does vector search\n",
    "        .select([\"text\", \"id__\"])  # retrieve the `text` and `id__` columns\n",
    "        .limit(4)  # get the top 4 rows\n",
    "        .to_list()\n",
    "    )\n",
    "\n",
    "    return state.update(relevant_chunks=search_results).append(chat_history=user_query)\n",
    "\n",
    "\n",
    "@action(reads=[\"chat_history\", \"relevant_chunks\"], writes=[\"chat_history\"])\n",
    "def bot_turn(state: State, llm_client: openai.OpenAI) -> State:\n",
    "    \"\"\"Collect relevant chunks and produce a response to the user query\"\"\"\n",
    "    user_query = state[\"chat_history\"][-1]\n",
    "    relevant_chunks = state[\"relevant_chunks\"]\n",
    "\n",
    "    # create system and user prompts\n",
    "    system_prompt = textwrap.dedent(\n",
    "        \"\"\"You are a conversational agent designed to discuss and provide \\\n",
    "        insights about various blog posts. Your task is to engage users in \\\n",
    "        meaningful conversations based on the content of the blog articles they mention.\n",
    "        \"\"\"\n",
    "    )\n",
    "    joined_chunks = ' '.join([c[\"text\"] for c in relevant_chunks])\n",
    "    user_prompt = \"BLOGS CONTENT\\n\" + joined_chunks + \"\\nUSER QUERY\\n\" + user_query\n",
    "\n",
    "    # query the OpenAI API\n",
    "    response = llm_client.chat.completions.create(\n",
    "        model=\"gpt-4o-mini\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt}\n",
    "        ]\n",
    "    )\n",
    "    bot_answer = response.choices[0].message.content\n",
    "\n",
    "    return state.append(chat_history=bot_answer)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Assemble the `Application`\n",
    "To build a Burr `Application`, you need to pass it actions and define valid transitions as tuples  `(from, to)`. The application must also define an `entrypoint` from where to begin execution. Then, we can visualize the graph of possible states and actions.\n",
    "\n",
    "First, let's see the simplest `Application` definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.43.0 (0)\n",
       " -->\n",
       "<!-- Title: %3 Pages: 1 -->\n",
       "<svg width=\"417pt\" height=\"177pt\"\n",
       " viewBox=\"0.00 0.00 416.50 177.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 173)\">\n",
       "<title>%3</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-173 412.5,-173 412.5,4 -4,4\"/>\n",
       "<!-- relevant_chunk_retrieval -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>relevant_chunk_retrieval</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M331.5,-103C331.5,-103 163.5,-103 163.5,-103 157.5,-103 151.5,-97 151.5,-91 151.5,-91 151.5,-78 151.5,-78 151.5,-72 157.5,-66 163.5,-66 163.5,-66 331.5,-66 331.5,-66 337.5,-66 343.5,-72 343.5,-78 343.5,-78 343.5,-91 343.5,-91 343.5,-97 337.5,-103 331.5,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"247.5\" y=\"-80.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">relevant_chunk_retrieval</text>\n",
       "</g>\n",
       "<!-- bot_turn -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>bot_turn</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M275.5,-37C275.5,-37 219.5,-37 219.5,-37 213.5,-37 207.5,-31 207.5,-25 207.5,-25 207.5,-12 207.5,-12 207.5,-6 213.5,0 219.5,0 219.5,0 275.5,0 275.5,0 281.5,0 287.5,-6 287.5,-12 287.5,-12 287.5,-25 287.5,-25 287.5,-31 281.5,-37 275.5,-37\"/>\n",
       "<text text-anchor=\"middle\" x=\"247.5\" y=\"-14.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">bot_turn</text>\n",
       "</g>\n",
       "<!-- relevant_chunk_retrieval&#45;&gt;bot_turn -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>relevant_chunk_retrieval&#45;&gt;bot_turn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M241.38,-65.67C240.84,-59.99 240.65,-53.55 240.8,-47.33\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"244.3,-47.42 241.39,-37.23 237.31,-47.01 244.3,-47.42\"/>\n",
       "</g>\n",
       "<!-- input__user_query -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__user_query</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"236.5,-169 90.5,-169 90.5,-132 236.5,-132 236.5,-169\"/>\n",
       "<text text-anchor=\"middle\" x=\"163.5\" y=\"-146.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: user_query</text>\n",
       "</g>\n",
       "<!-- input__user_query&#45;&gt;relevant_chunk_retrieval -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__user_query&#45;&gt;relevant_chunk_retrieval</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M186.9,-131.67C196.08,-124.68 206.74,-116.56 216.54,-109.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"218.68,-111.86 224.52,-103.01 214.44,-106.29 218.68,-111.86\"/>\n",
       "</g>\n",
       "<!-- input__lancedb_con -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>input__lancedb_con</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"408.5,-169 254.5,-169 254.5,-132 408.5,-132 408.5,-169\"/>\n",
       "<text text-anchor=\"middle\" x=\"331.5\" y=\"-146.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: lancedb_con</text>\n",
       "</g>\n",
       "<!-- input__lancedb_con&#45;&gt;relevant_chunk_retrieval -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>input__lancedb_con&#45;&gt;relevant_chunk_retrieval</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M308.1,-131.67C298.92,-124.68 288.26,-116.56 278.46,-109.09\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"280.56,-106.29 270.48,-103.01 276.32,-111.86 280.56,-106.29\"/>\n",
       "</g>\n",
       "<!-- bot_turn&#45;&gt;relevant_chunk_retrieval -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>bot_turn&#45;&gt;relevant_chunk_retrieval</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M253.61,-37.23C254.16,-42.91 254.35,-49.34 254.2,-55.57\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"250.7,-55.49 253.62,-65.67 257.69,-55.89 250.7,-55.49\"/>\n",
       "</g>\n",
       "<!-- input__llm_client -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>input__llm_client</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"133,-103 0,-103 0,-66 133,-66 133,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"66.5\" y=\"-80.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: llm_client</text>\n",
       "</g>\n",
       "<!-- input__llm_client&#45;&gt;bot_turn -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>input__llm_client&#45;&gt;bot_turn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M116.44,-65.84C141.86,-56.85 172.67,-45.96 198.03,-36.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"199.19,-40.29 207.45,-33.66 196.86,-33.69 199.19,-40.29\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7fbe85a88ed0>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from burr.core import ApplicationBuilder\n",
    "\n",
    "application = (\n",
    "    ApplicationBuilder()\n",
    "    .with_actions(relevant_chunk_retrieval, bot_turn)\n",
    "    .with_transitions(\n",
    "        (\"relevant_chunk_retrieval\", \"bot_turn\"),\n",
    "        (\"bot_turn\", \"relevant_chunk_retrieval\"),\n",
    "    )\n",
    "    .with_entrypoint(\"relevant_chunk_retrieval\")\n",
    "    .build()\n",
    ")\n",
    "application.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `ApplicationBuilder` patterns allows you to add all the features you need for production-readiness without modifying the logic of your agent. In the next few cells we'll add:\n",
    "\n",
    "- a hook to display the bot replies\n",
    "- tracking and storing execution metadata\n",
    "- add OpenTelemetry support\n",
    "\n",
    "To quickly develop an interactive experience in the terminal, we can add a `Hook` that will run after each `step` (i.e., `action`). It will check if the name of the previously completed action is equal to `bot_turn`. If it's the case, the hook prints the most recent message from the state's `chat_history`, in other words, the bot's reply."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from burr.lifecycle import PostRunStepHook\n",
    "\n",
    "class PrintBotAnswer(PostRunStepHook):\n",
    "    \"\"\"Hook to print the bot's answer\"\"\"\n",
    "    def post_run_step(self, *, state, action, **future_kwargs):\n",
    "        if action.name == \"bot_turn\":\n",
    "            print(\"\\n🤖: \", state[\"chat_history\"][-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, we specify via `.with_hooks()` to use the `PrintBotAnswer` hook and via `.with_tracker(..., use_otel_tracing=True)` to track execution and activate OpenTelemetry support. By importing the `opentelemetry.instrumentation` packages for `openai` and `lancedb` and using their `Instrumentor`, we'll be able to track more execution metadata.\n",
    "\n",
    "Also, we create the `OpenAI` client and the `LanceDBConnection` and bind to specific action parameter using the `.bind()` parameter. Notice how it changes the visualization slightly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.43.0 (0)\n",
       " -->\n",
       "<!-- Title: %3 Pages: 1 -->\n",
       "<svg width=\"200pt\" height=\"177pt\"\n",
       " viewBox=\"0.00 0.00 200.00 177.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 173)\">\n",
       "<title>%3</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-173 196,-173 196,4 -4,4\"/>\n",
       "<!-- relevant_chunk_retrieval -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>relevant_chunk_retrieval</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M180,-103C180,-103 12,-103 12,-103 6,-103 0,-97 0,-91 0,-91 0,-78 0,-78 0,-72 6,-66 12,-66 12,-66 180,-66 180,-66 186,-66 192,-72 192,-78 192,-78 192,-91 192,-91 192,-97 186,-103 180,-103\"/>\n",
       "<text text-anchor=\"middle\" x=\"96\" y=\"-80.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">relevant_chunk_retrieval</text>\n",
       "</g>\n",
       "<!-- bot_turn -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>bot_turn</title>\n",
       "<path fill=\"#b4d8e4\" stroke=\"black\" d=\"M124,-37C124,-37 68,-37 68,-37 62,-37 56,-31 56,-25 56,-25 56,-12 56,-12 56,-6 62,0 68,0 68,0 124,0 124,0 130,0 136,-6 136,-12 136,-12 136,-25 136,-25 136,-31 130,-37 124,-37\"/>\n",
       "<text text-anchor=\"middle\" x=\"96\" y=\"-14.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">bot_turn</text>\n",
       "</g>\n",
       "<!-- relevant_chunk_retrieval&#45;&gt;bot_turn -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>relevant_chunk_retrieval&#45;&gt;bot_turn</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M89.88,-65.67C89.34,-59.99 89.15,-53.55 89.3,-47.33\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"92.8,-47.42 89.89,-37.23 85.81,-47.01 92.8,-47.42\"/>\n",
       "</g>\n",
       "<!-- input__user_query -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>input__user_query</title>\n",
       "<polygon fill=\"none\" stroke=\"black\" stroke-dasharray=\"5,2\" points=\"169,-169 23,-169 23,-132 169,-132 169,-169\"/>\n",
       "<text text-anchor=\"middle\" x=\"96\" y=\"-146.8\" font-family=\"Helvetica,sans-Serif\" font-size=\"14.00\">input: user_query</text>\n",
       "</g>\n",
       "<!-- input__user_query&#45;&gt;relevant_chunk_retrieval -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>input__user_query&#45;&gt;relevant_chunk_retrieval</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M96,-131.67C96,-125.99 96,-119.55 96,-113.33\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"99.5,-113.23 96,-103.23 92.5,-113.23 99.5,-113.23\"/>\n",
       "</g>\n",
       "<!-- bot_turn&#45;&gt;relevant_chunk_retrieval -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>bot_turn&#45;&gt;relevant_chunk_retrieval</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M102.11,-37.23C102.66,-42.91 102.85,-49.34 102.7,-55.57\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"99.2,-55.49 102.12,-65.67 106.19,-55.89 99.2,-55.49\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7fbecfcbe5d0>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from burr.core import ApplicationBuilder\n",
    "from burr.integrations.opentelemetry import init_instruments\n",
    "\n",
    "init_instruments(\"openai\", \"lancedb\")\n",
    "\n",
    "\n",
    "llm_client = openai.OpenAI()\n",
    "lancedb_con = lancedb.connect(os.environ[\"DESTINATION__LANCEDB__CREDENTIALS__URI\"])\n",
    "\n",
    "application = (\n",
    "    ApplicationBuilder()\n",
    "    .with_actions(\n",
    "        relevant_chunk_retrieval.bind(lancedb_con=lancedb_con),\n",
    "        bot_turn.bind(llm_client=llm_client),\n",
    "    )\n",
    "    .with_transitions(\n",
    "        (\"relevant_chunk_retrieval\", \"bot_turn\"),\n",
    "        (\"bot_turn\", \"relevant_chunk_retrieval\"),\n",
    "    )\n",
    "    .with_entrypoint(\"relevant_chunk_retrieval\")\n",
    "    .with_tracker(\"local\", project=\"substack-rag\", use_otel_tracing=True)\n",
    "    .with_hooks(PrintBotAnswer())\n",
    "    .build()\n",
    ")\n",
    "application.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Launch the `Application`\n",
    "Finally, we start the application in a `while` loop, allowing it to run until we exit by inputting `quit` or `q`. We use `application.run()` and specify to halt after the action `bot_turn` to wait for the user's input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "## Lauching RAG application ##\n",
      "\n",
      "🤖:  Burr can be incredibly helpful if you're developing applications that require state management and debugging capabilities. Here are a few ways it can assist you:\n",
      "\n",
      "1. **State-Driven Logic**: Burr allows you to structure your application around actions that depend on the state, giving you fine control over the flow of your application. This is particularly useful for complex applications where the next step depends on the current state.\n",
      "\n",
      "2. **Easy Debugging**: With Burr, you can gain visibility into the decisions your application makes. If something goes wrong, you can rewind to a previous state and inspect what happened leading up to the issue. This makes identifying bugs much easier.\n",
      "\n",
      "3. **Forking Applications**: The ability to fork your application from any point in time allows you to create different branches of your application state, which can be useful for testing alternate paths or scenarios.\n",
      "\n",
      "4. **Integration with Other Frameworks**: Burr is designed to complement other frameworks, such as Hamilton. This makes it versatile and allows you to integrate it into your existing projects seamlessly.\n",
      "\n",
      "5. **User Input Handling**: You can define how your application interacts with user inputs, enabling the creation of interactive applications that respond dynamically to user actions.\n",
      "\n",
      "6. **Simplified Action Declarations**: By defining actions as functions or objects, you can keep your code organized and easy to follow. This can lead to more maintainable code, especially in large projects.\n",
      "\n",
      "If you're working on applications that require robust state management, dynamic responses, and efficient debugging, Burr could be a fantastic addition to your toolkit. Is there a specific type of application you have in mind that you'd like to develop with Burr?\n"
     ]
    }
   ],
   "source": [
    "# Launch the Burr application in a `while` loop\n",
    "print(\"\\n## Lauching RAG application ##\")\n",
    "user_query = input(\"\\nAsk something or type `quit/q` to exit: \")\n",
    "\n",
    "while True:\n",
    "    if user_query.lower() in [\"quit\", \"q\"]:\n",
    "        break\n",
    "\n",
    "    _, _, _ = application.run(\n",
    "        halt_after=[\"bot_turn\"],\n",
    "        inputs={\"user_query\": user_query},\n",
    "    )\n",
    "    user_query = input(\"\\nAsk something or type `quit/q` to exit: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Explore Burr UI\n",
    "After using the `Application`, use the command `burr` to launch the Burr UI.\n",
    "\n",
    "Burr UI allows you to:\n",
    "- explore past executions and current ones in real-time\n",
    "- see the application move through states\n",
    "- view code code failures\n",
    "- inspect tracked attributes (LLM prompts and responses, vector database calls, token counts)\n",
    "- one-click to create test fixtures from specific states\n",
    "- and more!\n",
    "\n",
    "![](burr-ui.gif)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
